package com.hup.utils.commons.thread;

import com.hup.utils.commons.exception.MultiException;
import lombok.AccessLevel;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.Setter;
import lombok.extern.log4j.Log4j2;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.util.Assert;

import java.util.ArrayList;
import java.util.Vector;
import java.util.concurrent.ThreadPoolExecutor.AbortPolicy;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

/**
 *
 * @author hugan
 * @date 2019/8/14
 */
@Log4j2
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public class MultiThreadExecutor {

    /**
     * 构造时固定的线程数(包含)
     */
    private int threadCountMax;
    private ThreadPoolTaskExecutor executor;
    /**
     * 任务总数
     */
    private int totalTask;
    /**
     * 使用线程数(<=threadCountFinal)
     */
    private int threadCount;
    /**
     * 单线程处理任务数(小于等于)
     */
    private int splitTask;

    private ProgressInfo progressInfo = new ProgressInfo();
    private boolean showInitLog = false;

    @Override
    public String toString() {
        return "MultiThreadExecutor{" +
                "totalTask=" + totalTask +
                ", threadCount=" + threadCount +
                ", splitTask=" + splitTask +
                '}';
    }

    private void initFixSplit(int totalTask) {
        Assert.isTrue(totalTask > 0, "任务数量为:" + totalTask);
        this.totalTask = totalTask;
        threadCount = totalTask / splitTask;
        if (totalTask % splitTask != 0) {
            threadCount += 1;
        }
        log.debug("{}", this);
    }

    /**
     * @param totalTask 本次任务数量
     * @param threadCount 本次指定线程数
     */
    private void initFixThread(int totalTask, int threadCount) {
        Assert.isTrue(totalTask > 0, "任务数量为:" + totalTask);
        Assert.isTrue(threadCount > 0, "任务数量为:" + threadCount);
        this.totalTask = totalTask;
        threadCount = Math.min(threadCount, threadCountMax);

        if (totalTask <= threadCount) {
            this.threadCount = totalTask;
            this.splitTask = 1;
        } else {
            this.threadCount = threadCount;
            this.splitTask = totalTask / this.threadCount;
            if (totalTask % this.threadCount != 0) {
                this.splitTask += 1;
                    /*
                    splitTask加了1后,可能出现实际线程数<threadCount的情况
                    [4,3][5,4][113,112]...
                     */
                int realThreadCount = totalTask / this.splitTask;
                if (realThreadCount < this.threadCount) {
                    this.threadCount = realThreadCount;
                    if (totalTask % this.splitTask != 0) this.threadCount += 1;
                }
            }
        }
        if (showInitLog) {
            log.info("{}", this);
        }
    }

    private Vector<Exception> iterator(TaskHandler taskHandler) {
        int lastTask = 0;
        progressInfo.reset();
        Vector<Exception> lstErr = new Vector<>();
        for (int tIndex = 0; tIndex < threadCount; tIndex++) {
            int stTask = splitTask * tIndex + 1;//起始于1
            int edTask = Math.min(stTask + splitTask - 1, totalTask);

            if (stTask > totalTask) {
                String msg = "解析异常,任务数越界! stTask=" + stTask + ", info=" + this;
                lstErr.add(new RuntimeException(msg));
                return lstErr;
            }
            lastTask = edTask;

            final int threadIndex = tIndex + 1;
            //taskHandler.onHandle(threadIndex, stTask, edTask);//原始线程按序执行
            //executor.execute(() -> taskHandler.onHandle(threadIndex, stTask, edTask, executor));//test
            /*
            add是应该在[主线程]的,要和waitFinish是一齐的
            若在初始化时就定了ThreadCount,则tIndex遍历异常时,可能会导致progressInfo死锁,
             */
            progressInfo.addTask();
            executor.execute(() -> {
                SplitInfo splitInfo = new SplitInfo(threadIndex, threadCount, stTask, edTask);
                try {
                    taskHandler.onHandle(splitInfo);
                } catch (Exception e) {
                    boolean handled = taskHandler.onException(splitInfo, e);
                    if (!handled) lstErr.add(e);
                } finally {
                    progressInfo.finishTask();
                }
            });
        }
        if (lastTask < totalTask) {
            String msg = "解析异常,任务数遗漏! lastTask=" + lastTask + ", info=" + this;
            lstErr.add(new RuntimeException(msg));
        }
        log.debug("executor={}", executor);

        progressInfo.waitFinish();
        return lstErr;
    }

    public void run(int totalTask, TaskHandler handler) {
        initFixThread(totalTask, threadCountMax);
        Vector<Exception> exceptions = iterator(handler);
        MultiException.doAssert(new ArrayList<>(exceptions));
    }

    public void showInitLog(boolean b) {
        showInitLog = b;
    }

    public void run(int totalTask, int threadCount, TaskHandler handler) {
        if (totalTask != 0) {
            initFixThread(totalTask, threadCount);
            Vector<Exception> exceptions = iterator(handler);
            MultiException.doAssert(new ArrayList<>(exceptions));
        }
    }

    public static MultiThreadExecutor fixThread(String threadNamePrefix, int maxThread) {
        return fixThread(threadNamePrefix, maxThread, 5);
    }

    public static MultiThreadExecutor fixThread(String threadNamePrefix, int maxThread, int keepAliveSeconds) {
        Assert.isTrue(maxThread > 0, "线程数量为:" + maxThread);

        MultiThreadExecutor res = new MultiThreadExecutor();
        res.threadCountMax = maxThread;

        res.executor = new ThreadPoolTaskExecutor();
        res.executor.setCorePoolSize(0);//核心线程池数量
        res.executor.setMaxPoolSize(maxThread);
        res.executor.setQueueCapacity(0);
        res.executor.setThreadNamePrefix(threadNamePrefix);
        res.executor.setRejectedExecutionHandler(new AbortPolicy());//超出任务数时抛异常,不会发生的
        res.executor.setKeepAliveSeconds(keepAliveSeconds);
        res.executor.initialize();
        return res;
    }

    public static MultiThreadExecutor fixSplit(String threadNamePrefix, int splitTask) {
        Assert.isTrue(splitTask > 0, "单线程任务数为:" + splitTask);
        throw new IllegalArgumentException("未实现");
    }

    /**
     * 用于处理多线程任务运行时,等待全部线程结束后才return
     * 处理线程等待唤醒
     * CountDownLatch是原生的,但不能复用
     */
    private static class ProgressInfo {
        /**
         * 当前真正运行的任务线程数
         */
        private int runningThread;
        private final ReentrantLock lock = new ReentrantLock(true);
        private final Condition condition = lock.newCondition();

        private void reset() {
            runningThread = 0;
        }

        private void addTask() {
            log.debug("st.runningThread={}", runningThread);
            lock.lock();
            log.debug("lock.in.before={}", runningThread);
            runningThread++;
            log.debug("lock.in.after={}", runningThread);
            lock.unlock();
            log.debug("ed.runningThread={}", runningThread);
        }

        private void finishTask() {
            log.debug("st.runningThread={}", runningThread);
            lock.lock();
            log.debug("lock.in.before={}", runningThread);
            runningThread--;
            log.debug("lock.in.after={}", runningThread);
            condition.signal();//唤醒waitFinish中等待的线程
            lock.unlock();
            log.debug("ed.runningThread={}", runningThread);
        }

        private void waitFinish() {
            log.debug("st.runningThread={}", runningThread);
            lock.lock();
            try {
                log.debug("while.st, runningThread={}", runningThread);
                while (runningThread != 0) {
                    log.debug("while.await, runningThread={}", runningThread);
                    condition.await();
                }
                log.debug("while.out, runningThread={}", runningThread);
            } catch (InterruptedException e) {
                log.error("InterruptedException", e);
                Thread.currentThread().interrupt();
            } finally {
                lock.unlock();
            }
        }
    }

    /**
     * 任务处理接口
     */
    public interface TaskHandler {
        void onHandle(SplitInfo splitInfo) throws Exception;

        /**
         * 处理子线程中发送的异常
         * @return handled:true-重写方法后应返回true,代表[已处理],适合无限期轮询的(等待用户取消的)任务;
         * false-省缺状态,所有子线程中的异常都会在任务结束后被抛出 {@link MultiException},适合一般耗时的(但可执行完的)任务
         */
        default boolean onException(SplitInfo splitInfo, Exception e) {
            return false;
        }
    }

    /**
     * 片段任务信息
     */
    public static class SplitInfo {
        /**
         * 当前是第n个线程,起始于1
         */
        public final int threadIndex;
        /**
         * 总线程数
         */
        public final int threadCount;
        /**
         * 当前线程的开始任务index(包含),起始于1
         */
        public final int stTask;
        /**
         * 当前线程的结束任务index(包含),起始于1
         */
        public final int edTask;
        /**
         * 任务数量=edTask-stTask+1
         */
        public final int splitTaskCount;

        /**
         * split自定义描述,可用为异常时的描述
         */
        @Setter
        @Getter
        private String desc;

        private SplitInfo(int threadIndex, int threadCount, int stTask, int edTask) {
            this.threadIndex = threadIndex;
            this.threadCount = threadCount;
            this.stTask = stTask;
            this.edTask = edTask;
            this.splitTaskCount = edTask - stTask + 1;
        }

        @Override
        public String toString() {
            return "SplitInfo{" +
                    threadIndex + "/" + threadCount +
                    " - [" + stTask + ", " + edTask + "]=" + splitTaskCount + "}";
        }
    }
}
